#include "arm_compute/runtime/NEON/NEFunctions.h"
#include "arm_compute/runtime/Tensor.h"
#include "arm_compute/runtime/TensorAllocator.h"
#include "arm_compute/core/TensorShape.h"
#include "arm_compute/core/Types.h"
#include "utils/Utils.h"

using namespace arm_compute;
using namespace utils;


#include <chrono>


/**
 * @brief 获取时间计数，单位为秒
 *
 * @return 秒
 * @note 在linux上返回的自1970-1-1以来经过的秒数，
 *  在windows上返回的是自机器最近一次启动以来经过的秒数
 */
inline double GetTime(void) {
  return std::chrono::duration<double>(
      std::chrono::high_resolution_clock::now().time_since_epoch()).count();
}

#ifdef __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
#pragma message "__ARM_FEATURE_FP16_VECTOR_ARITHMETIC defined"
#else
#pragma message "__ARM_FEATURE_FP16_VECTOR_ARITHMETIC not defined"
#endif

int main()
{
    // 矩阵维度
    const int M = 1000;
    const int N = 1000;
    const int K = 1000;

    // 创建张量
    Tensor A, B, C;
    A.allocator()->init(TensorInfo(TensorShape(K, M), 1, DataType::F16));
    B.allocator()->init(TensorInfo(TensorShape(N, K), 1, DataType::F16));
    C.allocator()->init(TensorInfo(TensorShape(N, M), 1, DataType::F16));

    // 分配张量内存
    A.allocator()->allocate();
    B.allocator()->allocate();
    C.allocator()->allocate();

    // 填充张量数据
    {
        Window window;
        window.use_tensor_dimensions(A.info()->tensor_shape());
        Iterator itA(&A, window);
        execute_window_loop(window, [&](const Coordinates & id)
        {
            *reinterpret_cast<float16_t *>(itA.ptr()) = static_cast<float16_t >(id.y() * K + id.x() + 1);
        }, itA);

        window.use_tensor_dimensions(B.info()->tensor_shape());
        Iterator itB(&B, window);
        execute_window_loop(window, [&](const Coordinates & id)
        {
            *reinterpret_cast<float16_t  *>(itB.ptr()) = static_cast<float16_t >(id.y() * N + id.x() + 1);
        }, itB);

        window.use_tensor_dimensions(C.info()->tensor_shape());
        Iterator itC(&C, window);
        execute_window_loop(window, [&](const Coordinates & id)
        {
            *reinterpret_cast<float16_t  *>(itC.ptr()) = 0.0f;
        }, itC);
    }

    // 创建矩阵乘法和加法函数
    NEGEMM gemm;
    gemm.configure(&A, &B, nullptr, &C, 1.0f, 0.0f);


    auto start_time = GetTime();
    for (int i = 0; i < 1000; i++)
    {
        // 运行函数
        gemm.run();
    }
    auto end_time = GetTime();
    std::cout<< end_time - start_time << "s" << std::endl;

    return 0;
}
